import cv2
import numpy as np
from PIL import ImageFont, ImageDraw, Image
import random

# 获取TFT显示屏中的ROI
def getMaxRect(img):
    # 将原图像灰度化
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 平均滤波
    blur = cv2.blur(gray, (5, 5))
    # 简单阈值的二值化
    # ret, thresh = cv.threshold(blur, 0,255,cv.THRESH_BINARY | cv.THRESH_TRIANGLE)
    ret, thresh = cv2.threshold(blur, 170, 255, cv2.THRESH_BINARY)

    # cv2.imshow('out1', thresh)

    # 查找轮廓
    contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, \
                                          cv2.CHAIN_APPROX_SIMPLE)[-2:]
    print(len(contours))
    # 找最大面积轮廓
    area_max = 0
    _cnt = None
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        area = cv2.contourArea(cnt)
        if x < 5 or y < 5 or (gray.shape[1] - 5 <= x <= gray.shape[1]) or (gray.shape[0] - 5 <= y <= gray.shape[0]):
            print("过滤边框", x, y, w, h)
            continue
        if area > area_max:
            area_max = area
            _cnt = cnt
    # 直边界矩形
    x, y, w, h = cv2.boundingRect(_cnt)
    print(x, y, w, h)
    if len(contours) > 0 and x > 0 and y > 0:
        img = img[y:h + y, x:w + x]

    return img

def putText(img, text, org=(0, 0), font_path='./simsun.ttc',
            color=(0, 0, 255), font_size = 30):
    """
    在图片上显示文字
    :param img: 输入的img, 通过cv2读取
    :param text: 要显示的文字
    :param org: 文字左上角坐标
    :param font_path: 字体路径
    :param color: 字体颜色, (B,G,R)
    :return:
    """
    img_pil = Image.fromarray(img)
    draw = ImageDraw.Draw(img_pil)
    b, g, r = color
    a = 0
    draw.text(org, text, font=ImageFont.truetype(font_path, font_size), fill=(b, g, r, a))
    img = np.array(img_pil)
    return img

def softmax(x):
    x -= np.max(x, axis=1, keepdims=True)  # 为了稳定地计算softmax概率， 一般会减掉最大的那个元素
    # print("减去行最大值 ：\n", x)
    x = np.exp(x) / np.sum(np.exp(x), axis=1, keepdims=True)
    return x

def decode_bbox(pre, priors, variances=None):
    if variances is None:
        variances = [0.1, 0.2]
    centers = priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:]

    sides = priors[:, 2:] * np.exp(pre[:, 2:] * variances[1])
    return np.concatenate([centers - sides / 2, centers + sides / 2], axis=1)


def showImage(img, boxes, classes, scores, img_height, img_width, prior_index, class_list):
    """
    draw bboxes and labels
    out:boxes,classes,scores
    """
    # bbox
    x1, y1, x2, y2 = int(boxes[prior_index][0] * img_width), int(boxes[prior_index][1] * img_height), \
                     int(boxes[prior_index][2] * img_width), int(boxes[prior_index][3] * img_height)
    if classes[prior_index] == 1:
        color = (0, 255, 0)
    else:
        color = (0, 0, 255)
    cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
    # confidence

    score = "{:.4f}".format(scores[prior_index])
    class_name = class_list[classes[prior_index]]

    cv2.putText(img, '{} {}'.format(class_name, score),
                (int(boxes[prior_index][0] * img_width), int(boxes[prior_index][1] * img_height) - 4),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255))


def showImageZH(img, boxes, classes, scores, img_height, img_width, prior_index, class_list):
    """
    draw bboxes and labels
    out:boxes,classes,scores
    """
    # bbox
    x1, y1, x2, y2 = int(boxes[prior_index][0] * img_width), int(boxes[prior_index][1] * img_height), \
                     int(boxes[prior_index][2] * img_width), int(boxes[prior_index][3] * img_height)
    # if classes[prior_index] == 1:
    #     color = (0, 255, 0)
    # else:
    #     color = (0, 0, 255)

    color = (255, 0, 255)

    img_copy = img.copy()
    cv2.rectangle(img_copy, (x1, y1), (x2, y2), color, -1)
    # confidence
    # cv2.rectangle(img_copy, (x1, y1), (x2, y2), (255, 0, 0), thickness=-1)
    cv2.addWeighted(img_copy, 0.2, img, 0.8, 0, img)
    class_name = class_list[classes[prior_index]]

    img = putText(img, class_name,
                  (int(boxes[prior_index][0] * img_width), int(boxes[prior_index][1] * img_height)-30),
                  color=(0, 0, 255))
    return img


def nms(box, scor, thresh):
    # 首先数据赋值和计算对应矩形框的面积
    # dets的数据格式是dets[[xmin,ymin,xmax,ymax,scores]....]
    x1 = box[:, 0]
    y1 = box[:, 1]
    x2 = box[:, 2]
    y2 = box[:, 3]
    areas = (y2 - y1 + 1) * (x2 - x1 + 1)
    scores = scor
    # print('areas  ', areas)
    # print('scores ', scores)

    # 这边的keep用于存放，NMS后剩余的方框
    keep = []

    # 取出分数从大到小排列的索引。.argsort()是从小到大排列，[::-1]是列表头和尾颠倒一下。
    index = scores.argsort()[::-1]
    # print(index)
    # 上面这两句比如分数[0.72 0.8  0.92 0.72 0.81 0.9 ]
    #  对应的索引index[  2   5    4     1    3   0  ]记住是取出索引，scores列表没变。

    # index会剔除遍历过的方框，和合并过的方框。
    while index.size > 0:
        # print(index.size)
        # 取出第一个方框进行和其他方框比对，看有没有可以合并的
        i = index[0]  # every time the first is the biggst, and add it directly

        # 因为我们这边分数已经按从大到小排列了。
        # 所以如果有合并存在，也是保留分数最高的这个，也就是我们现在那个这个
        # keep保留的是索引值，不是具体的分数。
        keep.append(i)
        # print(keep)
        # print('x1', x1[i])
        # print(x1[index[1:]])

        # 计算交集的左上角和右下角
        # 这里要注意，比如x1[i]这个方框的左上角x和所有其他的方框的左上角x的
        x11 = np.maximum(x1[i], x1[index[1:]])  # calculate the points of overlap
        y11 = np.maximum(y1[i], y1[index[1:]])
        x22 = np.minimum(x2[i], x2[index[1:]])
        y22 = np.minimum(y2[i], y2[index[1:]])

        # print(x11, y11, x22, y22)
        # 这边要注意，如果两个方框相交，X22-X11和Y22-Y11是正的。
        # 如果两个方框不相交，X22-X11和Y22-Y11是负的，我们把不相交的W和H设为0.
        w = np.maximum(0, x22 - x11 + 1)
        h = np.maximum(0, y22 - y11 + 1)

        # 计算重叠面积就是上面说的交集面积。不相交因为W和H都是0，所以不相交面积为0
        overlaps = w * h
        # print('overlaps is', overlaps)

        # 这个就是IOU公式（交并比）。
        # 得出来的ious是一个列表，里面拥有当前方框和其他所有方框的IOU结果。
        ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
        # print('ious is', ious)

        # 接下来是合并重叠度最大的方框，也就是合并ious中值大于thresh的方框
        # 我们合并的操作就是把他们剔除，因为我们合并这些方框只保留下分数最高的。
        # 我们经过排序当前我们操作的方框就是分数最高的，所以我们剔除其他和当前重叠度最高的方框
        # 这里np.where(ious<=thresh)[0]是一个固定写法。
        idx = np.where(ious <= thresh)[0]
        # print(idx)

        # 把留下来框在进行NMS操作
        # 这边留下的框是去除当前操作的框，和当前操作的框重叠度大于thresh的框
        # 每一次都会先去除当前操作框，所以索引的列表就会向前移动移位，要还原就+1，向后移动一位
        index = index[idx + 1]  # because index start from 1
        # print(index)
    return keep


def parsePredict(predictions, priors, cfg):
    label_classes = cfg['labels_list']
    # print('predictions[0]', predictions[0].shape)
    bbox_regressions, confs = np.split(predictions[0], [4, ], axis=1)
    boxes = decode_bbox(bbox_regressions, priors, cfg['variances'])

    confs = softmax(confs)

    out_boxes = []
    out_labels = []
    out_scores = []

    for c in range(1, len(label_classes)):
        cls = confs[:, c]
        score_idx = cls > cfg['score_threshold']
        cls_boxes = boxes[score_idx]
        cls_scores = cls[score_idx]

        # print('cls_boxes', cls_boxes)

        nms_idx = nms(cls_boxes, cls_scores, 0.8)
        # print("nms_idx", type(nms_idx), nms_idx)

        cls_boxes = cls_boxes[nms_idx]
        cls_scores = cls_scores[nms_idx]

        # print("cls_boxes", cls_boxes.shape)
        # print("cls_scores", cls_scores.shape)

        cls_labels = [c] * cls_boxes.shape[0]

        out_boxes.append(cls_boxes)
        out_labels.extend(cls_labels)
        out_scores.append(cls_scores)

    boxes = np.concatenate(out_boxes, axis=0)
    scores = np.concatenate(out_scores, axis=0)
    classes = np.array(out_labels)
    return boxes, classes, scores

